{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e5c6594-e5e8-4966-8cb8-a3e6a9ed7d89",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp common._base_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fce0c950-2e03-4be1-95d4-a02409d8dba3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c7c2ba5-19ee-421e-9252-7224b03f5201",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import inspect\n",
    "import random\n",
    "import warnings\n",
    "from contextlib import contextmanager\n",
    "from copy import deepcopy\n",
    "from dataclasses import dataclass\n",
    "from typing import List, Dict, Union\n",
    "\n",
    "import fsspec\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import pytorch_lightning as pl\n",
    "import neuralforecast.losses.pytorch as losses\n",
    "\n",
    "from neuralforecast.losses.pytorch import BasePointLoss, DistributionLoss\n",
    "from pytorch_lightning.callbacks.early_stopping import EarlyStopping\n",
    "from neuralforecast.tsdataset import (\n",
    "    TimeSeriesDataModule,\n",
    "    BaseTimeSeriesDataset,\n",
    "    _DistributedTimeSeriesDataModule,\n",
    ")\n",
    "from neuralforecast.common._scalers import TemporalNorm\n",
    "from neuralforecast.utils import get_indexer_raise_missing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6d4c4fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@dataclass\n",
    "class DistributedConfig:\n",
    "    partitions_path: str\n",
    "    num_nodes: int\n",
    "    devices: int"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5197e340-11f1-4c8c-96d1-ed396ac2b710",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "@contextmanager\n",
    "def _disable_torch_init():\n",
    "    \"\"\"Context manager used to disable pytorch's weight initialization.\n",
    "\n",
    "    This is especially useful when loading saved models, since when initializing\n",
    "    a model the weights are also initialized following some method\n",
    "    (e.g. kaiming uniform), and that time is wasted since we'll override them with\n",
    "    the saved weights.\"\"\"\n",
    "    def noop(*args, **kwargs):\n",
    "        return\n",
    "        \n",
    "    kaiming_uniform = nn.init.kaiming_uniform_\n",
    "    kaiming_normal = nn.init.kaiming_normal_\n",
    "    xavier_uniform = nn.init.xavier_uniform_\n",
    "    xavier_normal = nn.init.xavier_normal_\n",
    "    \n",
    "    nn.init.kaiming_uniform_ = noop\n",
    "    nn.init.kaiming_normal_ = noop\n",
    "    nn.init.xavier_uniform_ = noop\n",
    "    nn.init.xavier_normal_ = noop\n",
    "    try:\n",
    "        yield\n",
    "    finally:\n",
    "        nn.init.kaiming_uniform_ = kaiming_uniform\n",
    "        nn.init.kaiming_normal_ = kaiming_normal\n",
    "        nn.init.xavier_uniform_ = xavier_uniform\n",
    "        nn.init.xavier_normal_ = xavier_normal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f58c0e89",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "def tensor_to_numpy(tensor: torch.Tensor) -> np.ndarray:\n",
    "    \"\"\"Convert a tensor to numpy\"\"\"\n",
    "    if tensor.dtype == torch.bfloat16:\n",
    "        return tensor.float().numpy()\n",
    "    \n",
    "    return tensor.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60c40a64-8381-46a2-8cbb-70ec70ed7914",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class BaseModel(pl.LightningModule):\n",
    "    EXOGENOUS_FUTR = True   # If the model can handle future exogenous variables\n",
    "    EXOGENOUS_HIST = True   # If the model can handle historical exogenous variables\n",
    "    EXOGENOUS_STAT = True   # If the model can handle static exogenous variables\n",
    "    MULTIVARIATE = False    # If the model produces multivariate forecasts (True) or univariate (False)\n",
    "    RECURRENT = False       # If the model produces forecasts recursively (True) or direct (False)\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        h: int,\n",
    "        input_size: int,\n",
    "        loss: Union[BasePointLoss, DistributionLoss, nn.Module],\n",
    "        valid_loss: Union[BasePointLoss, DistributionLoss, nn.Module],\n",
    "        learning_rate: float,\n",
    "        max_steps: int,\n",
    "        val_check_steps: int,\n",
    "        batch_size: int,\n",
    "        valid_batch_size: Union[int, None],\n",
    "        windows_batch_size: int,\n",
    "        inference_windows_batch_size: Union[int, None],\n",
    "        start_padding_enabled: bool,\n",
    "        n_series: Union[int, None] = None,\n",
    "        n_samples: Union[int, None] = 100,\n",
    "        h_train: int = 1,\n",
    "        inference_input_size: Union[int, None] = None,\n",
    "        step_size: int = 1,\n",
    "        num_lr_decays: int = 0,\n",
    "        early_stop_patience_steps: int = -1,\n",
    "        scaler_type: str = 'identity',\n",
    "        futr_exog_list: Union[List, None] = None,\n",
    "        hist_exog_list: Union[List, None] = None,\n",
    "        stat_exog_list: Union[List, None] = None,\n",
    "        exclude_insample_y: Union[bool, None] = False,\n",
    "        drop_last_loader: Union[bool, None] = False,\n",
    "        random_seed: Union[int, None] = 1,\n",
    "        alias: Union[str, None] = None,\n",
    "        optimizer: Union[torch.optim.Optimizer, None] = None,\n",
    "        optimizer_kwargs: Union[Dict, None] = None,\n",
    "        lr_scheduler: Union[torch.optim.lr_scheduler.LRScheduler, None] = None,\n",
    "        lr_scheduler_kwargs: Union[Dict, None] = None,\n",
    "        dataloader_kwargs=None,\n",
    "        **trainer_kwargs,\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        # Multivarariate checks\n",
    "        if self.MULTIVARIATE and n_series is None:\n",
    "            raise Exception(f'{type(self).__name__} is a multivariate model. Please set n_series to the number of unique time series in your dataset.')\n",
    "        if not self.MULTIVARIATE:\n",
    "            n_series = 1\n",
    "        self.n_series = n_series          \n",
    "\n",
    "        # Protections for previous recurrent models\n",
    "        if input_size < 1:\n",
    "            input_size = 3 * h\n",
    "            warnings.warn(\n",
    "                f'Input size too small. Automatically setting input size to 3 * horizon = {input_size}'\n",
    "            )\n",
    "\n",
    "        if inference_input_size is None:\n",
    "            inference_input_size = input_size            \n",
    "        elif inference_input_size is not None and inference_input_size < 1:\n",
    "            inference_input_size = input_size\n",
    "            warnings.warn(\n",
    "                f'Inference input size too small. Automatically setting inference input size to input_size = {input_size}'\n",
    "            )\n",
    "\n",
    "        # For recurrent models we need one additional input as we need to shift insample_y to use it as input\n",
    "        if self.RECURRENT:\n",
    "            input_size += 1\n",
    "            inference_input_size += 1\n",
    "\n",
    "        # Attributes needed for recurrent models\n",
    "        self.horizon_backup = h\n",
    "        self.input_size_backup = input_size\n",
    "        self.n_samples = n_samples\n",
    "        if self.RECURRENT:\n",
    "            if (\n",
    "                hasattr(loss, \"horizon_weight\")\n",
    "                and loss.horizon_weight is not None\n",
    "                and h_train != h\n",
    "            ):\n",
    "                warnings.warn(\n",
    "                    f\"Setting h_train={h} to match the horizon_weight length.\"\n",
    "                )\n",
    "                h_train = h\n",
    "            self.h_train = h_train\n",
    "            self.inference_input_size = inference_input_size\n",
    "            self.rnn_state = None\n",
    "            self.maintain_state = False\n",
    "\n",
    "        with warnings.catch_warnings(record=False):\n",
    "            warnings.filterwarnings('ignore')\n",
    "            # the following line issues a warning about the loss attribute being saved\n",
    "            # but we do want to save it\n",
    "            self.save_hyperparameters() # Allows instantiation from a checkpoint from class\n",
    "        self.random_seed = random_seed\n",
    "        pl.seed_everything(self.random_seed, workers=True)\n",
    "\n",
    "        # Loss\n",
    "        self.loss = loss\n",
    "        if valid_loss is None:\n",
    "            self.valid_loss = loss\n",
    "        else:\n",
    "            self.valid_loss = valid_loss\n",
    "        self.train_trajectories: List = []\n",
    "        self.valid_trajectories: List = []\n",
    "\n",
    "        # Optimization\n",
    "        if optimizer is not None and not issubclass(optimizer, torch.optim.Optimizer):\n",
    "            raise TypeError(\"optimizer is not a valid subclass of torch.optim.Optimizer\")\n",
    "        self.optimizer = optimizer\n",
    "        self.optimizer_kwargs = optimizer_kwargs if optimizer_kwargs is not None else {}\n",
    "\n",
    "        # lr scheduler\n",
    "        if lr_scheduler is not None and not issubclass(lr_scheduler, torch.optim.lr_scheduler.LRScheduler):\n",
    "            raise TypeError(\"lr_scheduler is not a valid subclass of torch.optim.lr_scheduler.LRScheduler\")\n",
    "        self.lr_scheduler = lr_scheduler\n",
    "        self.lr_scheduler_kwargs = lr_scheduler_kwargs if lr_scheduler_kwargs is not None else {}\n",
    "\n",
    "        # Variables\n",
    "        self.futr_exog_list = list(futr_exog_list) if futr_exog_list is not None else []\n",
    "        self.hist_exog_list = list(hist_exog_list) if hist_exog_list is not None else []\n",
    "        self.stat_exog_list = list(stat_exog_list) if stat_exog_list is not None else []\n",
    "\n",
    "        # Set data sizes\n",
    "        self.futr_exog_size = len(self.futr_exog_list)\n",
    "        self.hist_exog_size = len(self.hist_exog_list)\n",
    "        self.stat_exog_size = len(self.stat_exog_list)   \n",
    "\n",
    "        # Check if model supports exogenous, otherwise raise Exception\n",
    "        if not self.EXOGENOUS_FUTR and self.futr_exog_size > 0:\n",
    "            raise Exception(f'{type(self).__name__} does not support future exogenous variables.')\n",
    "        if not self.EXOGENOUS_HIST and self.hist_exog_size > 0:\n",
    "            raise Exception(f'{type(self).__name__} does not support historical exogenous variables.')\n",
    "        if not self.EXOGENOUS_STAT and self.stat_exog_size > 0:\n",
    "            raise Exception(f'{type(self).__name__} does not support static exogenous variables.')\n",
    "\n",
    "        # Protections for loss functions\n",
    "        if isinstance(self.loss, (losses.IQLoss, losses.HuberIQLoss)):\n",
    "            loss_type = type(self.loss)\n",
    "            if not isinstance(self.valid_loss, loss_type):\n",
    "                raise Exception(f'Please set valid_loss={type(self.loss).__name__}() when training with {type(self.loss).__name__}')\n",
    "        if isinstance(self.loss, (losses.MQLoss, losses.HuberMQLoss)):\n",
    "            if not isinstance(self.valid_loss, (losses.MQLoss, losses.HuberMQLoss)):\n",
    "                raise Exception(f'Please set valid_loss to MQLoss() or HuberMQLoss() when training with {type(self.loss).__name__}')\n",
    "        if isinstance(self.valid_loss, (losses.IQLoss, losses.HuberIQLoss)):\n",
    "            valid_loss_type = type(self.valid_loss)\n",
    "            if not isinstance(self.loss, valid_loss_type):\n",
    "                raise Exception(f'Please set loss={type(self.valid_loss).__name__}() when validating with {type(self.valid_loss).__name__}')        \n",
    "\n",
    "        # Deny impossible loss / valid_loss combinations\n",
    "        if isinstance(self.loss, losses.BasePointLoss) and self.valid_loss.is_distribution_output:\n",
    "            raise Exception(f'Validation with distribution loss {type(self.valid_loss).__name__} is not possible when using loss={type(self.loss).__name__}. Please use a point valid_loss (MAE, MSE, ...)')\n",
    "        elif self.valid_loss.is_distribution_output and self.valid_loss is not loss:\n",
    "            # Maybe we should raise a Warning or an Exception here, but meh for now.\n",
    "            self.valid_loss = loss\n",
    "        \n",
    "        if isinstance(self.loss, (losses.relMSE, losses.Accuracy, losses.sCRPS)):\n",
    "            raise Exception(f\"{type(self.loss).__name__} cannot be used for training. Please use another loss function (MAE, MSE, ...)\")\n",
    "        \n",
    "        if isinstance(self.valid_loss, (losses.relMSE)):\n",
    "            raise Exception(f\"{type(self.valid_loss).__name__} cannot be used for validation. Please use another valid_loss (MAE, MSE, ...)\")\n",
    "\n",
    "        ## Trainer arguments ##\n",
    "        # Max steps, validation steps and check_val_every_n_epoch\n",
    "        trainer_kwargs = {**trainer_kwargs, 'max_steps': max_steps}\n",
    "\n",
    "        if 'max_epochs' in trainer_kwargs.keys():\n",
    "            raise Exception('max_epochs is deprecated, use max_steps instead.')\n",
    "\n",
    "        # Callbacks\n",
    "        if early_stop_patience_steps > 0:\n",
    "            if 'callbacks' not in trainer_kwargs:\n",
    "                trainer_kwargs['callbacks'] = []\n",
    "            trainer_kwargs['callbacks'].append(\n",
    "                EarlyStopping(\n",
    "                    monitor='ptl/val_loss', patience=early_stop_patience_steps\n",
    "                )\n",
    "            )\n",
    "\n",
    "        # Add GPU accelerator if available\n",
    "        if trainer_kwargs.get('accelerator', None) is None:\n",
    "            if torch.cuda.is_available():\n",
    "                trainer_kwargs['accelerator'] = \"gpu\"\n",
    "        if trainer_kwargs.get('devices', None) is None:\n",
    "            if torch.cuda.is_available():\n",
    "                trainer_kwargs['devices'] = -1\n",
    "\n",
    "        # Avoid saturating local memory, disabled fit model checkpoints\n",
    "        if trainer_kwargs.get('enable_checkpointing', None) is None:\n",
    "            trainer_kwargs['enable_checkpointing'] = False\n",
    "\n",
    "        # Set other attributes\n",
    "        self.trainer_kwargs = trainer_kwargs\n",
    "        self.h = h\n",
    "        self.input_size = input_size\n",
    "        self.windows_batch_size = windows_batch_size\n",
    "        self.start_padding_enabled = start_padding_enabled\n",
    "\n",
    "        # Padder to complete train windows, \n",
    "        # example y=[1,2,3,4,5] h=3 -> last y_output = [5,0,0]\n",
    "        if start_padding_enabled:\n",
    "            self.padder_train = nn.ConstantPad1d(padding=(self.input_size-1, self.h), value=0.0)\n",
    "        else:\n",
    "            self.padder_train = nn.ConstantPad1d(padding=(0, self.h), value=0.0)\n",
    "\n",
    "        # Batch sizes\n",
    "        if self.MULTIVARIATE and n_series is not None:\n",
    "            self.batch_size = max(batch_size, n_series)\n",
    "            if valid_batch_size is not None:\n",
    "                valid_batch_size = max(valid_batch_size, n_series)\n",
    "        else:\n",
    "            self.batch_size = batch_size\n",
    "        \n",
    "        if valid_batch_size is None:\n",
    "            self.valid_batch_size = self.batch_size\n",
    "        else:\n",
    "            self.valid_batch_size = valid_batch_size\n",
    "\n",
    "        if inference_windows_batch_size is None:\n",
    "            self.inference_windows_batch_size = windows_batch_size\n",
    "        else:\n",
    "            self.inference_windows_batch_size = inference_windows_batch_size\n",
    "\n",
    "        # Optimization \n",
    "        self.learning_rate = learning_rate\n",
    "        self.max_steps = max_steps\n",
    "        self.num_lr_decays = num_lr_decays\n",
    "        self.lr_decay_steps = (\n",
    "            max(max_steps // self.num_lr_decays, 1) if self.num_lr_decays > 0 else 10e7\n",
    "        )\n",
    "        self.early_stop_patience_steps = early_stop_patience_steps\n",
    "        self.val_check_steps = val_check_steps\n",
    "        self.windows_batch_size = windows_batch_size\n",
    "        self.step_size = step_size\n",
    "        \n",
    "        # If the model does not support exogenous, it can't support exclude_insample_y\n",
    "        if exclude_insample_y and not (self.EXOGENOUS_FUTR or self.EXOGENOUS_HIST or self.EXOGENOUS_STAT):\n",
    "            raise Exception(f'{type(self).__name__} does not support `exclude_insample_y=True`. Please set `exclude_insample_y=False`')\n",
    "\n",
    "        self.exclude_insample_y = exclude_insample_y\n",
    "\n",
    "        # Scaler\n",
    "        self.scaler = TemporalNorm(\n",
    "            scaler_type=scaler_type,\n",
    "            dim=1,  # Time dimension is 1.\n",
    "            num_features= 1 + len(self.hist_exog_list) + len(self.futr_exog_list)\n",
    "        )\n",
    "\n",
    "        # Fit arguments\n",
    "        self.val_size = 0\n",
    "        self.test_size = 0\n",
    "\n",
    "        # Model state\n",
    "        self.decompose_forecast = False\n",
    "\n",
    "        # DataModule arguments\n",
    "        self.dataloader_kwargs = dataloader_kwargs\n",
    "        self.drop_last_loader = drop_last_loader\n",
    "        # used by on_validation_epoch_end hook\n",
    "        self.validation_step_outputs: List = []\n",
    "        self.alias = alias\n",
    "\n",
    "    def __repr__(self):\n",
    "        return type(self).__name__ if self.alias is None else self.alias\n",
    "\n",
    "    def _check_exog(self, dataset):\n",
    "        temporal_cols = set(dataset.temporal_cols.tolist())\n",
    "        static_cols = set(dataset.static_cols.tolist() if dataset.static_cols is not None else [])\n",
    "\n",
    "        missing_hist = set(self.hist_exog_list) - temporal_cols\n",
    "        missing_futr = set(self.futr_exog_list) - temporal_cols\n",
    "        missing_stat = set(self.stat_exog_list) - static_cols\n",
    "        if missing_hist:\n",
    "            raise Exception(f'{missing_hist} historical exogenous variables not found in input dataset')\n",
    "        if missing_futr:\n",
    "            raise Exception(f'{missing_futr} future exogenous variables not found in input dataset')\n",
    "        if missing_stat:\n",
    "            raise Exception(f'{missing_stat} static exogenous variables not found in input dataset')\n",
    "\n",
    "    def _restart_seed(self, random_seed):\n",
    "        if random_seed is None:\n",
    "            random_seed = self.random_seed\n",
    "        torch.manual_seed(random_seed)\n",
    "\n",
    "    def _get_temporal_exogenous_cols(self, temporal_cols):\n",
    "        return list(\n",
    "            set(temporal_cols.tolist()) & set(self.hist_exog_list + self.futr_exog_list)\n",
    "        )\n",
    "    \n",
    "    def _set_quantiles(self, quantiles=None):\n",
    "        if quantiles is None and isinstance(self.loss, (losses.IQLoss, losses.HuberIQLoss)):\n",
    "            self.loss.update_quantile(q=[0.5])\n",
    "        elif hasattr(self.loss, 'update_quantile') and callable(self.loss.update_quantile):\n",
    "            self.loss.update_quantile(q=quantiles)\n",
    "\n",
    "    def _fit_distributed(\n",
    "        self,\n",
    "        distributed_config,\n",
    "        datamodule,\n",
    "        val_size,\n",
    "        test_size,\n",
    "    ):\n",
    "        assert distributed_config is not None\n",
    "        from pyspark.ml.torch.distributor import TorchDistributor\n",
    "\n",
    "        def train_fn(\n",
    "            model_cls,\n",
    "            model_params,\n",
    "            datamodule,\n",
    "            trainer_kwargs,\n",
    "            num_tasks,\n",
    "            num_proc_per_task,\n",
    "            val_size,\n",
    "            test_size,\n",
    "        ):\n",
    "            import pytorch_lightning as pl\n",
    "\n",
    "            # we instantiate here to avoid pickling large tensors (weights)\n",
    "            model = model_cls(**model_params)\n",
    "            model.val_size = val_size\n",
    "            model.test_size = test_size\n",
    "            for arg in ('devices', 'num_nodes'):\n",
    "                trainer_kwargs.pop(arg, None)\n",
    "            trainer = pl.Trainer(\n",
    "                strategy=\"ddp\",\n",
    "                use_distributed_sampler=False,  # to ensure our dataloaders are used as-is\n",
    "                num_nodes=num_tasks,\n",
    "                devices=num_proc_per_task,\n",
    "                **trainer_kwargs,\n",
    "            )\n",
    "            trainer.fit(model=model, datamodule=datamodule)\n",
    "            model.metrics = trainer.callback_metrics\n",
    "            model.__dict__.pop('_trainer', None)\n",
    "            return model\n",
    "\n",
    "        def is_gpu_accelerator(accelerator):\n",
    "            from pytorch_lightning.accelerators.cuda import CUDAAccelerator\n",
    "\n",
    "            return (\n",
    "                accelerator == \"gpu\"\n",
    "                or isinstance(accelerator, CUDAAccelerator)\n",
    "                or (accelerator == \"auto\" and CUDAAccelerator.is_available())\n",
    "            )\n",
    "\n",
    "        local_mode = distributed_config.num_nodes == 1\n",
    "        if local_mode:\n",
    "            num_tasks = 1\n",
    "            num_proc_per_task = distributed_config.devices\n",
    "        else:\n",
    "            num_tasks = distributed_config.num_nodes * distributed_config.devices\n",
    "            num_proc_per_task = 1  # number of GPUs per task\n",
    "        num_proc = num_tasks * num_proc_per_task\n",
    "        use_gpu = is_gpu_accelerator(self.trainer_kwargs[\"accelerator\"])\n",
    "        model = TorchDistributor(\n",
    "            num_processes=num_proc,\n",
    "            local_mode=local_mode,\n",
    "            use_gpu=use_gpu,\n",
    "        ).run(\n",
    "            train_fn,\n",
    "            model_cls=type(self),\n",
    "            model_params=self.hparams,\n",
    "            datamodule=datamodule,\n",
    "            trainer_kwargs=self.trainer_kwargs,\n",
    "            num_tasks=num_tasks,\n",
    "            num_proc_per_task=num_proc_per_task,\n",
    "            val_size=val_size,\n",
    "            test_size=test_size,\n",
    "        )\n",
    "        return model\n",
    "\n",
    "    def _fit(\n",
    "        self,\n",
    "        dataset,\n",
    "        batch_size,\n",
    "        valid_batch_size=1024,\n",
    "        val_size=0,\n",
    "        test_size=0,\n",
    "        random_seed=None,\n",
    "        shuffle_train=True,\n",
    "        distributed_config=None,\n",
    "    ):\n",
    "        self._check_exog(dataset)\n",
    "        self._restart_seed(random_seed)\n",
    "\n",
    "        self.val_size = val_size\n",
    "        self.test_size = test_size\n",
    "        is_local = isinstance(dataset, BaseTimeSeriesDataset)\n",
    "        if is_local:\n",
    "            datamodule_constructor = TimeSeriesDataModule\n",
    "        else:\n",
    "            datamodule_constructor = _DistributedTimeSeriesDataModule\n",
    "        \n",
    "        dataloader_kwargs = self.dataloader_kwargs if self.dataloader_kwargs is not None else {}\n",
    "        datamodule = datamodule_constructor(\n",
    "            dataset=dataset, \n",
    "            batch_size=batch_size,\n",
    "            valid_batch_size=valid_batch_size,\n",
    "            drop_last=self.drop_last_loader,\n",
    "            shuffle_train=shuffle_train,\n",
    "            **dataloader_kwargs\n",
    "        )\n",
    "\n",
    "        if self.val_check_steps > self.max_steps:\n",
    "            warnings.warn(\n",
    "                'val_check_steps is greater than max_steps, '\n",
    "                'setting val_check_steps to max_steps.'\n",
    "            )\n",
    "        val_check_interval = min(self.val_check_steps, self.max_steps)\n",
    "        self.trainer_kwargs['val_check_interval'] = int(val_check_interval)\n",
    "        self.trainer_kwargs['check_val_every_n_epoch'] = None\n",
    "\n",
    "        if is_local:\n",
    "            model = self\n",
    "            trainer = pl.Trainer(**model.trainer_kwargs)\n",
    "            trainer.fit(model, datamodule=datamodule)\n",
    "            model.metrics = trainer.callback_metrics\n",
    "            model.__dict__.pop('_trainer', None)\n",
    "        else:\n",
    "            model = self._fit_distributed(\n",
    "                distributed_config,\n",
    "                datamodule,\n",
    "                val_size,\n",
    "                test_size,\n",
    "            )\n",
    "        return model\n",
    "\n",
    "    def on_fit_start(self):\n",
    "        torch.manual_seed(self.random_seed)\n",
    "        np.random.seed(self.random_seed)\n",
    "        random.seed(self.random_seed)\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        if self.optimizer:\n",
    "            optimizer_signature = inspect.signature(self.optimizer)\n",
    "            optimizer_kwargs = deepcopy(self.optimizer_kwargs)\n",
    "            if 'lr' in optimizer_signature.parameters:\n",
    "                if 'lr' in optimizer_kwargs:\n",
    "                    warnings.warn(\"ignoring learning rate passed in optimizer_kwargs, using the model's learning rate\")\n",
    "                optimizer_kwargs['lr'] = self.learning_rate\n",
    "            optimizer = self.optimizer(params=self.parameters(), **optimizer_kwargs)\n",
    "        else:\n",
    "            if self.optimizer_kwargs:\n",
    "                warnings.warn(\n",
    "                    \"ignoring optimizer_kwargs as the optimizer is not specified\"\n",
    "                )            \n",
    "            optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)\n",
    "        \n",
    "        lr_scheduler = {'frequency': 1, 'interval': 'step'}\n",
    "        if self.lr_scheduler:\n",
    "            lr_scheduler_signature = inspect.signature(self.lr_scheduler)\n",
    "            lr_scheduler_kwargs = deepcopy(self.lr_scheduler_kwargs)\n",
    "            if 'optimizer' in lr_scheduler_signature.parameters:\n",
    "                if 'optimizer' in lr_scheduler_kwargs:\n",
    "                    warnings.warn(\"ignoring optimizer passed in lr_scheduler_kwargs, using the model's optimizer\")\n",
    "                    del lr_scheduler_kwargs['optimizer']\n",
    "            lr_scheduler['scheduler'] = self.lr_scheduler(optimizer=optimizer, **lr_scheduler_kwargs)\n",
    "        else:\n",
    "            if self.lr_scheduler_kwargs:\n",
    "                warnings.warn(\n",
    "                    \"ignoring lr_scheduler_kwargs as the lr_scheduler is not specified\"\n",
    "                )            \n",
    "            lr_scheduler['scheduler'] = torch.optim.lr_scheduler.StepLR(\n",
    "                optimizer=optimizer, step_size=self.lr_decay_steps, gamma=0.5\n",
    "            )\n",
    "        return {'optimizer': optimizer, 'lr_scheduler': lr_scheduler}\n",
    "\n",
    "    def get_test_size(self):\n",
    "        return self.test_size\n",
    "\n",
    "    def set_test_size(self, test_size):\n",
    "        self.test_size = test_size\n",
    "\n",
    "    def on_validation_epoch_end(self):\n",
    "        if self.val_size == 0:\n",
    "            return\n",
    "        losses = torch.stack(self.validation_step_outputs)\n",
    "        avg_loss = losses.mean().detach().item()\n",
    "        self.log(\n",
    "            \"ptl/val_loss\",\n",
    "            avg_loss,\n",
    "            batch_size=losses.size(0),\n",
    "            sync_dist=True,\n",
    "        )\n",
    "        self.valid_trajectories.append((self.global_step, avg_loss))\n",
    "        self.validation_step_outputs.clear() # free memory (compute `avg_loss` per epoch)\n",
    "\n",
    "    def save(self, path):\n",
    "        with fsspec.open(path, 'wb') as f:\n",
    "            torch.save(\n",
    "                {'hyper_parameters': self.hparams, 'state_dict': self.state_dict()},\n",
    "                f,\n",
    "            )\n",
    "\n",
    "    @classmethod\n",
    "    def load(cls, path, **kwargs):\n",
    "        if \"weights_only\" in inspect.signature(torch.load).parameters:\n",
    "            kwargs[\"weights_only\"] = False\n",
    "        with fsspec.open(path, 'rb') as f, warnings.catch_warnings():\n",
    "            # ignore possible warnings about weights_only=False\n",
    "            warnings.filterwarnings('ignore', category=FutureWarning)\n",
    "            content = torch.load(f, **kwargs)\n",
    "        with _disable_torch_init():\n",
    "            model = cls(**content['hyper_parameters']) \n",
    "        if \"assign\" in inspect.signature(model.load_state_dict).parameters:\n",
    "            model.load_state_dict(content[\"state_dict\"], strict=True, assign=True)\n",
    "        else:  # pytorch<2.1\n",
    "            model.load_state_dict(content[\"state_dict\"], strict=True)\n",
    "        return model\n",
    "\n",
    "    def _create_windows(self, batch, step):\n",
    "        # Parse common data\n",
    "        window_size = self.input_size + self.h\n",
    "        temporal_cols = batch['temporal_cols']\n",
    "        temporal = batch['temporal']                \n",
    "\n",
    "        if step == 'train':\n",
    "            if self.val_size + self.test_size > 0:\n",
    "                cutoff = -self.val_size - self.test_size\n",
    "                temporal = temporal[:, :, :cutoff]\n",
    "\n",
    "            temporal = self.padder_train(temporal)\n",
    "            \n",
    "            if temporal.shape[-1] < window_size:\n",
    "                raise Exception('Time series is too short for training, consider setting a smaller input size or set start_padding_enabled=True')\n",
    "            \n",
    "            windows = temporal.unfold(dimension=-1, \n",
    "                                      size=window_size, \n",
    "                                      step=self.step_size)\n",
    "\n",
    "            if self.MULTIVARIATE:\n",
    "                # [n_series, C, Ws, L + h] -> [Ws, L + h, C, n_series]\n",
    "                windows = windows.permute(2, 3, 1, 0)\n",
    "            else:\n",
    "                # [n_series, C, Ws, L + h] -> [Ws * n_series, L + h, C, 1]\n",
    "                windows_per_serie = windows.shape[2]\n",
    "                windows = windows.permute(0, 2, 3, 1)\n",
    "                windows = windows.flatten(0, 1)\n",
    "                windows = windows.unsqueeze(-1)\n",
    "\n",
    "            # Sample and Available conditions\n",
    "            available_idx = temporal_cols.get_loc('available_mask')           \n",
    "            available_condition = windows[:, :self.input_size, available_idx]\n",
    "            available_condition = torch.sum(available_condition, axis=(1, -1)) # Sum over time & series dimension\n",
    "            final_condition = (available_condition > 0)\n",
    "            \n",
    "            if self.h > 0:\n",
    "                sample_condition = windows[:, self.input_size:, available_idx]\n",
    "                sample_condition = torch.sum(sample_condition, axis=(1, -1)) # Sum over time & series dimension\n",
    "                final_condition = (sample_condition > 0) & (available_condition > 0)\n",
    "            \n",
    "            windows = windows[final_condition]\n",
    "\n",
    "            # Parse Static data to match windows\n",
    "            static = batch.get(\"static\", None)\n",
    "            static_cols = batch.get(\"static_cols\", None)\n",
    "            \n",
    "            # Repeat static if univariate: [n_series, S] -> [Ws * n_series, S]\n",
    "            if static is not None and not self.MULTIVARIATE:\n",
    "                static = torch.repeat_interleave(static, \n",
    "                                    repeats=windows_per_serie, dim=0)\n",
    "                static = static[final_condition]        \n",
    "\n",
    "            # Protection of empty windows\n",
    "            if final_condition.sum() == 0:\n",
    "                raise Exception('No windows available for training')\n",
    "\n",
    "            return windows, static, static_cols\n",
    "\n",
    "        elif step in ['predict', 'val']:\n",
    "\n",
    "            if step == 'predict':\n",
    "                initial_input = temporal.shape[-1] - self.test_size\n",
    "                if initial_input <= self.input_size: # There is not enough data to predict first timestamp\n",
    "                    temporal = F.pad(temporal, pad=(self.input_size-initial_input, 0), mode=\"constant\", value=0.0)\n",
    "                predict_step_size = self.predict_step_size\n",
    "                cutoff = - self.input_size - self.test_size\n",
    "                temporal = temporal[:, :, cutoff:]\n",
    "\n",
    "            elif step == 'val':\n",
    "                predict_step_size = self.step_size\n",
    "                cutoff = -self.input_size - self.val_size - self.test_size\n",
    "                if self.test_size > 0:\n",
    "                    temporal = batch['temporal'][:, :, cutoff:-self.test_size]\n",
    "                else:\n",
    "                    temporal = batch['temporal'][:, :, cutoff:]\n",
    "                if temporal.shape[-1] < window_size:\n",
    "                    initial_input = temporal.shape[-1] - self.val_size\n",
    "                    temporal = F.pad(temporal, pad=(self.input_size-initial_input, 0), mode=\"constant\", value=0.0)\n",
    "\n",
    "            if (step=='predict') and (self.test_size==0) and (len(self.futr_exog_list)==0):\n",
    "                temporal = F.pad(temporal, pad=(0, self.h), mode=\"constant\", value=0.0)\n",
    "\n",
    "            windows = temporal.unfold(dimension=-1,\n",
    "                                      size=window_size,\n",
    "                                      step=predict_step_size)\n",
    "\n",
    "            static = batch.get('static', None)\n",
    "            static_cols=batch.get('static_cols', None)\n",
    "\n",
    "            if self.MULTIVARIATE:\n",
    "                # [n_series, C, Ws, L + h] -> [Ws, L + h, C, n_series]\n",
    "                windows = windows.permute(2, 3, 1, 0)\n",
    "            else:\n",
    "                # [n_series, C, Ws, L + h] -> [Ws * n_series, L + h, C, 1]\n",
    "                windows_per_serie = windows.shape[2]\n",
    "                windows = windows.permute(0, 2, 3, 1)\n",
    "                windows = windows.flatten(0, 1)\n",
    "                windows = windows.unsqueeze(-1)\n",
    "                if static is not None:\n",
    "                    static = torch.repeat_interleave(static, \n",
    "                                    repeats=windows_per_serie, dim=0)\n",
    "\n",
    "            return windows, static, static_cols\n",
    "        else:\n",
    "            raise ValueError(f'Unknown step {step}') \n",
    "\n",
    "    def _normalization(self, windows, y_idx):\n",
    "        # windows are already filtered by train/validation/test\n",
    "        # from the `create_windows_method` nor leakage risk\n",
    "        temporal = windows['temporal']                  # [Ws, L + h, C, n_series]\n",
    "        temporal_cols = windows['temporal_cols'].copy() # [Ws, L + h, C, n_series]\n",
    "\n",
    "        # To avoid leakage uses only the lags\n",
    "        temporal_data_cols = self._get_temporal_exogenous_cols(temporal_cols=temporal_cols)\n",
    "        temporal_idxs = get_indexer_raise_missing(temporal_cols, temporal_data_cols)\n",
    "        temporal_idxs = np.append(y_idx, temporal_idxs)\n",
    "        temporal_data = temporal[:, :, temporal_idxs] \n",
    "        temporal_mask = temporal[:, :, temporal_cols.get_loc('available_mask')].clone()\n",
    "        if self.h > 0:\n",
    "            temporal_mask[:, -self.h:] = 0.0\n",
    "\n",
    "        # Normalize. self.scaler stores the shift and scale for inverse transform\n",
    "        temporal_mask = temporal_mask.unsqueeze(2) # Add channel dimension for scaler.transform.\n",
    "        temporal_data = self.scaler.transform(x=temporal_data, mask=temporal_mask)\n",
    "\n",
    "        # Replace values in windows dict\n",
    "        temporal[:, :, temporal_idxs] = temporal_data\n",
    "        windows['temporal'] = temporal\n",
    "\n",
    "        return windows\n",
    "\n",
    "    def _inv_normalization(self, y_hat, y_idx):\n",
    "        # Receives window predictions [Ws, h, output, n_series]\n",
    "        # Broadcasts scale if necessary and inverts normalization\n",
    "        add_channel_dim = y_hat.ndim > 3\n",
    "        y_loc, y_scale = self._get_loc_scale(y_idx, add_channel_dim=add_channel_dim)\n",
    "        y_hat = self.scaler.inverse_transform(z=y_hat, x_scale=y_scale, x_shift=y_loc)\n",
    "\n",
    "        return y_hat\n",
    "    \n",
    "    def _sample_windows(self, windows_temporal, static, static_cols, temporal_cols, step, w_idxs=None):\n",
    "        if step == 'train' and self.windows_batch_size is not None:\n",
    "            n_windows = windows_temporal.shape[0]\n",
    "            w_idxs = np.random.choice(n_windows, \n",
    "                                        size=self.windows_batch_size,\n",
    "                                        replace=(n_windows < self.windows_batch_size))\n",
    "        windows_sample = windows_temporal\n",
    "        if w_idxs is not None:\n",
    "            windows_sample = windows_temporal[w_idxs]\n",
    "            \n",
    "            if static is not None and not self.MULTIVARIATE:\n",
    "                static = static[w_idxs]\n",
    "\n",
    "        windows_batch = dict(temporal=windows_sample,\n",
    "                                temporal_cols=temporal_cols,\n",
    "                                static=static,\n",
    "                                static_cols=static_cols)\n",
    "        return windows_batch\n",
    "\n",
    "    def _parse_windows(self, batch, windows):\n",
    "        # windows: [Ws, L + h, C, n_series]\n",
    "\n",
    "        # Filter insample lags from outsample horizon\n",
    "        y_idx = batch['y_idx']\n",
    "        mask_idx = batch['temporal_cols'].get_loc('available_mask')\n",
    "\n",
    "        insample_y = windows['temporal'][:, :self.input_size, y_idx]\n",
    "        insample_mask = windows['temporal'][:, :self.input_size, mask_idx]\n",
    "\n",
    "        # Declare additional information\n",
    "        outsample_y = None\n",
    "        outsample_mask = None\n",
    "        hist_exog = None\n",
    "        futr_exog = None\n",
    "        stat_exog = None\n",
    "\n",
    "        if self.h > 0:\n",
    "            outsample_y = windows['temporal'][:, self.input_size:, y_idx]\n",
    "            outsample_mask = windows['temporal'][:, self.input_size:, mask_idx]\n",
    "\n",
    "        # Recurrent models at t predict t+1, so we shift the input (insample_y) by one\n",
    "        if self.RECURRENT:\n",
    "            insample_y = torch.cat((insample_y, outsample_y[:, :-1]), dim=1)\n",
    "            insample_mask = torch.cat((insample_mask, outsample_mask[:, :-1]), dim=1)\n",
    "            self.maintain_state = False\n",
    "\n",
    "        if len(self.hist_exog_list):\n",
    "            hist_exog_idx = get_indexer_raise_missing(windows['temporal_cols'], self.hist_exog_list)\n",
    "            if self.RECURRENT:\n",
    "                hist_exog = windows['temporal'][:, :, hist_exog_idx]\n",
    "                hist_exog[:, self.input_size:] = 0.0\n",
    "                hist_exog = hist_exog[:, 1:]\n",
    "            else:\n",
    "                hist_exog = windows['temporal'][:, :self.input_size, hist_exog_idx]\n",
    "            if not self.MULTIVARIATE:\n",
    "                hist_exog = hist_exog.squeeze(-1)\n",
    "            else:\n",
    "                hist_exog = hist_exog.swapaxes(1, 2)\n",
    "\n",
    "        if len(self.futr_exog_list):\n",
    "            futr_exog_idx = get_indexer_raise_missing(windows['temporal_cols'], self.futr_exog_list)\n",
    "            futr_exog = windows['temporal'][:, :, futr_exog_idx]\n",
    "            if self.RECURRENT:\n",
    "                futr_exog = futr_exog[:, 1:]\n",
    "            if not self.MULTIVARIATE:\n",
    "                futr_exog = futr_exog.squeeze(-1)\n",
    "            else:\n",
    "                futr_exog = futr_exog.swapaxes(1, 2)                \n",
    "\n",
    "        if len(self.stat_exog_list):\n",
    "            static_idx = get_indexer_raise_missing(windows['static_cols'], self.stat_exog_list)\n",
    "            stat_exog = windows['static'][:, static_idx]\n",
    "\n",
    "        # TODO: think a better way of removing insample_y features\n",
    "        if self.exclude_insample_y:\n",
    "            insample_y = insample_y * 0\n",
    "\n",
    "        return insample_y, insample_mask, outsample_y, outsample_mask, \\\n",
    "               hist_exog, futr_exog, stat_exog     \n",
    "\n",
    "    def _get_loc_scale(self, y_idx, add_channel_dim=False):\n",
    "        # [B, L, C, n_series] -> [B, L, n_series]\n",
    "        y_scale = self.scaler.x_scale[:, :, y_idx]\n",
    "        y_loc = self.scaler.x_shift[:, :, y_idx]\n",
    "        \n",
    "        # [B, L, n_series] -> [B, L, n_series, 1]\n",
    "        if add_channel_dim:\n",
    "            y_scale = y_scale.unsqueeze(-1)\n",
    "            y_loc = y_loc.unsqueeze(-1)\n",
    "\n",
    "        return y_loc, y_scale\n",
    "\n",
    "    def _compute_valid_loss(self, insample_y, outsample_y, output, outsample_mask, y_idx):\n",
    "        if self.loss.is_distribution_output:\n",
    "            y_loc, y_scale = self._get_loc_scale(y_idx)\n",
    "            distr_args = self.loss.scale_decouple(output=output, loc=y_loc, scale=y_scale)\n",
    "            if isinstance(self.valid_loss, (losses.sCRPS, losses.MQLoss, losses.HuberMQLoss)):\n",
    "                _, _, quants  = self.loss.sample(distr_args=distr_args)            \n",
    "                output = quants\n",
    "            elif isinstance(self.valid_loss, losses.BasePointLoss):\n",
    "                distr = self.loss.get_distribution(distr_args=distr_args)\n",
    "                output = distr.mean\n",
    "\n",
    "        # Validation Loss evaluation\n",
    "        if self.valid_loss.is_distribution_output:\n",
    "            valid_loss = self.valid_loss(y=outsample_y, distr_args=distr_args, mask=outsample_mask)\n",
    "        else:\n",
    "            output = self._inv_normalization(y_hat=output, y_idx=y_idx)\n",
    "            valid_loss = self.valid_loss(y=outsample_y, y_hat=output, y_insample=insample_y, mask=outsample_mask)\n",
    "        return valid_loss\n",
    "    \n",
    "    def _validate_step_recurrent_batch(self, insample_y, insample_mask, futr_exog, hist_exog, stat_exog, y_idx):\n",
    "        # Remember state in network and set horizon to 1\n",
    "        self.rnn_state = None\n",
    "        self.maintain_state = True\n",
    "        self.h = 1\n",
    "\n",
    "        # Initialize results array\n",
    "        n_outputs = self.loss.outputsize_multiplier\n",
    "        y_hat = torch.zeros((insample_y.shape[0],\n",
    "                            self.horizon_backup,\n",
    "                            self.n_series * n_outputs),\n",
    "                            device=insample_y.device,\n",
    "                            dtype=insample_y.dtype)\n",
    "\n",
    "        # First step prediction\n",
    "        tau = 0\n",
    "        \n",
    "        # Set exogenous\n",
    "        hist_exog_current = None\n",
    "        if self.hist_exog_size > 0:\n",
    "            hist_exog_current = hist_exog[:, :self.input_size + tau]\n",
    "\n",
    "        futr_exog_current = None\n",
    "        if self.futr_exog_size > 0:\n",
    "            futr_exog_current = futr_exog[:, :self.input_size + tau]\n",
    "\n",
    "        # First forecast step\n",
    "        y_hat[:, tau], insample_y = self._validate_step_recurrent_single(\n",
    "                                                                insample_y=insample_y[:, :self.input_size + tau],\n",
    "                                                                insample_mask=insample_mask[:, :self.input_size + tau],\n",
    "                                                                hist_exog=hist_exog_current,\n",
    "                                                                futr_exog=futr_exog_current,\n",
    "                                                                stat_exog=stat_exog,\n",
    "                                                                y_idx=y_idx,\n",
    "                                                                )\n",
    "\n",
    "        # Horizon prediction recursively\n",
    "        for tau in range(1, self.horizon_backup):\n",
    "            # Set exogenous\n",
    "            if self.hist_exog_size > 0:\n",
    "                hist_exog_current = hist_exog[:, self.input_size + tau - 1].unsqueeze(1)\n",
    "\n",
    "            if self.futr_exog_size > 0:\n",
    "                futr_exog_current = futr_exog[:, self.input_size + tau - 1].unsqueeze(1)\n",
    "            \n",
    "            y_hat[:, tau], insample_y = self._validate_step_recurrent_single(\n",
    "                                                                insample_y=insample_y,\n",
    "                                                                insample_mask=None,\n",
    "                                                                hist_exog=hist_exog_current,\n",
    "                                                                futr_exog=futr_exog_current,\n",
    "                                                                stat_exog=stat_exog,\n",
    "                                                                y_idx = y_idx,\n",
    "                                                                )\n",
    "        \n",
    "        # Reset state and horizon\n",
    "        self.maintain_state = False\n",
    "        self.rnn_state = None\n",
    "        self.h = self.horizon_backup\n",
    "\n",
    "        return y_hat   \n",
    "\n",
    "    def _validate_step_recurrent_single(self, insample_y, insample_mask, hist_exog, futr_exog, stat_exog, y_idx):\n",
    "        # Input sequence\n",
    "        windows_batch = dict(insample_y=insample_y,                 # [Ws, L, n_series]\n",
    "                        insample_mask=insample_mask,                # [Ws, L, n_series]\n",
    "                        futr_exog=futr_exog,                        # univariate: [Ws, L, F]; multivariate: [Ws, F, L, n_series]\n",
    "                        hist_exog=hist_exog,                        # univariate: [Ws, L, X]; multivariate: [Ws, X, L, n_series]\n",
    "                        stat_exog=stat_exog)                        # univariate: [Ws, S]; multivariate: [n_series, S]\n",
    "\n",
    "        # Model Predictions\n",
    "        output_batch_unmapped = self(windows_batch)\n",
    "        output_batch = self.loss.domain_map(output_batch_unmapped)\n",
    "        \n",
    "        # Inverse normalization and sampling\n",
    "        if self.loss.is_distribution_output:\n",
    "            # Sample distribution\n",
    "            y_loc, y_scale = self._get_loc_scale(y_idx)\n",
    "            distr_args = self.loss.scale_decouple(output=output_batch, loc=y_loc, scale=y_scale)\n",
    "            # When validating, the output is the mean of the distribution which is an attribute\n",
    "            distr = self.loss.get_distribution(distr_args=distr_args)\n",
    "\n",
    "            # Scale back to feed back as input\n",
    "            insample_y = self.scaler.scaler(distr.mean, y_loc, y_scale)\n",
    "        else:\n",
    "            # Todo: for now, we assume that in case of a BasePointLoss with ndim==4, the last dimension\n",
    "            # contains a set of predictions for the target (e.g. MQLoss multiple quantiles), for which we use the \n",
    "            # mean as feedback signal for the recurrent predictions. A more precise way is to increase the\n",
    "            # insample input size of the recurrent network by the number of outputs so that each output\n",
    "            # can be fed back to a specific input channel. \n",
    "            if output_batch.ndim == 4:\n",
    "                output_batch = output_batch.mean(dim=-1)\n",
    "\n",
    "            insample_y = output_batch\n",
    "\n",
    "        # Remove horizon dim: [B, 1, N * n_outputs] -> [B, N * n_outputs]\n",
    "        y_hat = output_batch_unmapped.squeeze(1)\n",
    "        return y_hat, insample_y\n",
    "\n",
    "    def _predict_step_recurrent_batch(self, insample_y, insample_mask, futr_exog, hist_exog, stat_exog, y_idx):\n",
    "        # Remember state in network and set horizon to 1\n",
    "        self.rnn_state = None\n",
    "        self.maintain_state = True\n",
    "        self.h = 1\n",
    "\n",
    "        # Initialize results array\n",
    "        n_outputs = len(self.loss.output_names)\n",
    "        y_hat = torch.zeros((insample_y.shape[0],\n",
    "                            self.horizon_backup,\n",
    "                            self.n_series,\n",
    "                            n_outputs),\n",
    "                            device=insample_y.device,\n",
    "                            dtype=insample_y.dtype)\n",
    "\n",
    "        # First step prediction\n",
    "        tau = 0\n",
    "        \n",
    "        # Set exogenous\n",
    "        hist_exog_current = None\n",
    "        if self.hist_exog_size > 0:\n",
    "            hist_exog_current = hist_exog[:, :self.input_size + tau]\n",
    "\n",
    "        futr_exog_current = None\n",
    "        if self.futr_exog_size > 0:\n",
    "            futr_exog_current = futr_exog[:, :self.input_size + tau]\n",
    "\n",
    "        # First forecast step\n",
    "        y_hat[:, tau], insample_y = self._predict_step_recurrent_single(\n",
    "                                                                insample_y=insample_y[:, :self.input_size + tau],\n",
    "                                                                insample_mask=insample_mask[:, :self.input_size + tau],\n",
    "                                                                hist_exog=hist_exog_current,\n",
    "                                                                futr_exog=futr_exog_current,\n",
    "                                                                stat_exog=stat_exog,\n",
    "                                                                y_idx=y_idx,\n",
    "                                                                )\n",
    "\n",
    "        # Horizon prediction recursively\n",
    "        for tau in range(1, self.horizon_backup):\n",
    "            # Set exogenous\n",
    "            if self.hist_exog_size > 0:\n",
    "                hist_exog_current = hist_exog[:, self.input_size + tau - 1].unsqueeze(1)\n",
    "\n",
    "            if self.futr_exog_size > 0:\n",
    "                futr_exog_current = futr_exog[:, self.input_size + tau - 1].unsqueeze(1)\n",
    "            \n",
    "            y_hat[:, tau], insample_y = self._predict_step_recurrent_single(\n",
    "                                                                insample_y=insample_y,\n",
    "                                                                insample_mask=None,\n",
    "                                                                hist_exog=hist_exog_current,\n",
    "                                                                futr_exog=futr_exog_current,\n",
    "                                                                stat_exog=stat_exog,\n",
    "                                                                y_idx = y_idx,\n",
    "                                                                )\n",
    "        \n",
    "        # Reset state and horizon\n",
    "        self.maintain_state = False\n",
    "        self.rnn_state = None\n",
    "        self.h = self.horizon_backup\n",
    "\n",
    "        # Squeeze for univariate case\n",
    "        if not self.MULTIVARIATE:\n",
    "            y_hat = y_hat.squeeze(2)\n",
    "\n",
    "        return y_hat        \n",
    "\n",
    "    def _predict_step_recurrent_single(self, insample_y, insample_mask, hist_exog, futr_exog, stat_exog, y_idx):\n",
    "        # Input sequence\n",
    "        windows_batch = dict(insample_y=insample_y,                 # [Ws, L, n_series]\n",
    "                        insample_mask=insample_mask,                # [Ws, L, n_series]\n",
    "                        futr_exog=futr_exog,                        # univariate: [Ws, L, F]; multivariate: [Ws, F, L, n_series]\n",
    "                        hist_exog=hist_exog,                        # univariate: [Ws, L, X]; multivariate: [Ws, X, L, n_series]\n",
    "                        stat_exog=stat_exog)                        # univariate: [Ws, S]; multivariate: [n_series, S]\n",
    "\n",
    "        # Model Predictions\n",
    "        output_batch_unmapped = self(windows_batch)\n",
    "        output_batch = self.loss.domain_map(output_batch_unmapped)\n",
    "        \n",
    "        # Inverse normalization and sampling\n",
    "        if self.loss.is_distribution_output:\n",
    "            # Sample distribution\n",
    "            y_loc, y_scale = self._get_loc_scale(y_idx)\n",
    "            distr_args = self.loss.scale_decouple(output=output_batch, loc=y_loc, scale=y_scale)\n",
    "            # When predicting, we need to sample to get the quantiles. The mean is an attribute.\n",
    "            _, _, quants = self.loss.sample(distr_args=distr_args, num_samples=self.n_samples)\n",
    "            mean = self.loss.distr_mean\n",
    "\n",
    "            # Scale back to feed back as input\n",
    "            insample_y = self.scaler.scaler(mean, y_loc, y_scale)\n",
    "            \n",
    "            # Save predictions\n",
    "            y_hat = torch.concat((mean.unsqueeze(-1), quants), axis=-1)\n",
    "\n",
    "            if self.loss.return_params:\n",
    "                distr_args = torch.stack(distr_args, dim=-1)\n",
    "                if distr_args.ndim > 4:\n",
    "                    distr_args = distr_args.flatten(-2, -1)\n",
    "                y_hat = torch.concat((y_hat, distr_args), axis=-1)\n",
    "        else:\n",
    "            # Todo: for now, we assume that in case of a BasePointLoss with ndim==4, the last dimension\n",
    "            # contains a set of predictions for the target (e.g. MQLoss multiple quantiles), for which we use the \n",
    "            # mean as feedback signal for the recurrent predictions. A more precise way is to increase the\n",
    "            # insample input size of the recurrent network by the number of outputs so that each output\n",
    "            # can be fed back to a specific input channel. \n",
    "            if output_batch.ndim == 4:\n",
    "                output_batch = output_batch.mean(dim=-1)\n",
    "\n",
    "            insample_y = output_batch\n",
    "            y_hat = self._inv_normalization(y_hat=output_batch, y_idx=y_idx)\n",
    "            y_hat = y_hat.unsqueeze(-1)\n",
    "\n",
    "        # Remove horizon dim: [B, 1, N, n_outputs] -> [B, N, n_outputs]\n",
    "        y_hat = y_hat.squeeze(1)\n",
    "        return y_hat, insample_y\n",
    "\n",
    "    def _predict_step_direct_batch(self, insample_y, insample_mask, hist_exog, futr_exog, stat_exog, y_idx):\n",
    "        windows_batch = dict(insample_y=insample_y,                 # [Ws, L, n_series]\n",
    "                        insample_mask=insample_mask,                # [Ws, L, n_series]\n",
    "                        futr_exog=futr_exog,                        # univariate: [Ws, L, F]; multivariate: [Ws, F, L, n_series]\n",
    "                        hist_exog=hist_exog,                        # univariate: [Ws, L, X]; multivariate: [Ws, X, L, n_series]\n",
    "                        stat_exog=stat_exog)                        # univariate: [Ws, S]; multivariate: [n_series, S]\n",
    "\n",
    "        # Model Predictions\n",
    "        output_batch = self(windows_batch)\n",
    "        output_batch = self.loss.domain_map(output_batch)\n",
    "\n",
    "        # Inverse normalization and sampling\n",
    "        if self.loss.is_distribution_output:\n",
    "            y_loc, y_scale = self._get_loc_scale(y_idx)\n",
    "            distr_args = self.loss.scale_decouple(output=output_batch, loc=y_loc, scale=y_scale)\n",
    "            _, sample_mean, quants = self.loss.sample(distr_args=distr_args)\n",
    "            y_hat = torch.concat((sample_mean, quants), axis=-1)\n",
    "\n",
    "            if self.loss.return_params:\n",
    "                distr_args = torch.stack(distr_args, dim=-1)\n",
    "                if distr_args.ndim > 4:\n",
    "                    distr_args = distr_args.flatten(-2, -1)\n",
    "                y_hat = torch.concat((y_hat, distr_args), axis=-1)                \n",
    "        else:\n",
    "             y_hat = self._inv_normalization(y_hat=output_batch, \n",
    "                                            y_idx=y_idx)\n",
    "\n",
    "        return y_hat\n",
    "            \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        # Set horizon to h_train in case of recurrent model to speed up training\n",
    "        if self.RECURRENT:\n",
    "            self.h = self.h_train\n",
    "        \n",
    "        # windows: [Ws, L + h, C, n_series] or [Ws, L + h, C]\n",
    "        y_idx = batch['y_idx']\n",
    "\n",
    "        temporal_cols = batch['temporal_cols']\n",
    "        windows_temporal, static, static_cols = self._create_windows(batch, step='train')\n",
    "        windows = self._sample_windows(windows_temporal, static, static_cols, temporal_cols, step='train')\n",
    "        original_outsample_y = torch.clone(windows['temporal'][:, self.input_size:, y_idx])\n",
    "        windows = self._normalization(windows=windows, y_idx=y_idx)\n",
    "        \n",
    "        # Parse windows\n",
    "        insample_y, insample_mask, outsample_y, outsample_mask, \\\n",
    "               hist_exog, futr_exog, stat_exog = self._parse_windows(batch, windows)\n",
    "\n",
    "        windows_batch = dict(insample_y=insample_y,                 # [Ws, L, n_series]\n",
    "                        insample_mask=insample_mask,                # [Ws, L, n_series]\n",
    "                        futr_exog=futr_exog,                        # univariate: [Ws, L, F]; multivariate: [Ws, F, L, n_series]\n",
    "                        hist_exog=hist_exog,                        # univariate: [Ws, L, X]; multivariate: [Ws, X, L, n_series]\n",
    "                        stat_exog=stat_exog)                        # univariate: [Ws, S]; multivariate: [n_series, S]\n",
    "\n",
    "        # Model Predictions\n",
    "        output = self(windows_batch)\n",
    "        output = self.loss.domain_map(output)\n",
    "        \n",
    "        if self.loss.is_distribution_output:\n",
    "            y_loc, y_scale = self._get_loc_scale(y_idx)\n",
    "            outsample_y = original_outsample_y\n",
    "            distr_args = self.loss.scale_decouple(output=output, loc=y_loc, scale=y_scale)\n",
    "            loss = self.loss(y=outsample_y, distr_args=distr_args, mask=outsample_mask)\n",
    "        else:\n",
    "            loss = self.loss(y=outsample_y, y_hat=output, y_insample=insample_y, mask=outsample_mask)\n",
    "\n",
    "        if torch.isnan(loss):\n",
    "            print('Model Parameters', self.hparams)\n",
    "            print('insample_y', torch.isnan(insample_y).sum())\n",
    "            print('outsample_y', torch.isnan(outsample_y).sum())\n",
    "            raise Exception('Loss is NaN, training stopped.')\n",
    "\n",
    "        train_loss_log = loss.detach().item()\n",
    "        self.log(\n",
    "            'train_loss',\n",
    "            train_loss_log,\n",
    "            batch_size=outsample_y.size(0),\n",
    "            prog_bar=True,\n",
    "            on_epoch=True,\n",
    "        )\n",
    "        self.train_trajectories.append((self.global_step, train_loss_log))\n",
    "\n",
    "        self.h = self.horizon_backup\n",
    "\n",
    "        return loss\n",
    "\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        if self.val_size == 0:\n",
    "            return np.nan\n",
    "\n",
    "        temporal_cols = batch['temporal_cols']\n",
    "        windows_temporal, static, static_cols = self._create_windows(batch, step='val')\n",
    "        n_windows = len(windows_temporal)\n",
    "        y_idx = batch['y_idx']\n",
    "\n",
    "        # Number of windows in batch\n",
    "        windows_batch_size = self.inference_windows_batch_size\n",
    "        if windows_batch_size < 0:\n",
    "            windows_batch_size = n_windows\n",
    "        n_batches = int(np.ceil(n_windows / windows_batch_size))\n",
    "\n",
    "        valid_losses = []\n",
    "        batch_sizes = []\n",
    "        for i in range(n_batches):\n",
    "            # Create and normalize windows [Ws, L + h, C, n_series]\n",
    "            w_idxs = np.arange(i*windows_batch_size, \n",
    "                               min((i+1)*windows_batch_size, n_windows))\n",
    "            windows = self._sample_windows(windows_temporal, static, static_cols, temporal_cols, step='val', w_idxs=w_idxs)\n",
    "            original_outsample_y = torch.clone(windows['temporal'][:, self.input_size:, y_idx])\n",
    "\n",
    "            windows = self._normalization(windows=windows, y_idx=y_idx)\n",
    "\n",
    "            # Parse windows\n",
    "            insample_y, insample_mask, _, outsample_mask, \\\n",
    "                hist_exog, futr_exog, stat_exog = self._parse_windows(batch, windows)\n",
    "\n",
    "            if self.RECURRENT:\n",
    "                output_batch = self._validate_step_recurrent_batch(insample_y=insample_y,\n",
    "                                                           insample_mask=insample_mask,\n",
    "                                                           futr_exog=futr_exog,\n",
    "                                                           hist_exog=hist_exog,\n",
    "                                                           stat_exog=stat_exog,\n",
    "                                                           y_idx=y_idx)\n",
    "            else:       \n",
    "                windows_batch = dict(insample_y=insample_y,                 # [Ws, L, n_series]\n",
    "                                insample_mask=insample_mask,                # [Ws, L, n_series]\n",
    "                                futr_exog=futr_exog,                        # univariate: [Ws, L, F]; multivariate: [Ws, F, L, n_series]\n",
    "                                hist_exog=hist_exog,                        # univariate: [Ws, L, X]; multivariate: [Ws, X, L, n_series]\n",
    "                                stat_exog=stat_exog)                        # univariate: [Ws, S]; multivariate: [n_series, S]\n",
    "                \n",
    "                # Model Predictions\n",
    "                output_batch = self(windows_batch)   \n",
    "\n",
    "            output_batch = self.loss.domain_map(output_batch)\n",
    "            valid_loss_batch = self._compute_valid_loss(insample_y=insample_y,\n",
    "                                                        outsample_y=original_outsample_y,\n",
    "                                                output=output_batch, \n",
    "                                                outsample_mask=outsample_mask,\n",
    "                                                y_idx=batch['y_idx'])\n",
    "            valid_losses.append(valid_loss_batch)\n",
    "            batch_sizes.append(len(output_batch))\n",
    "        \n",
    "        valid_loss = torch.stack(valid_losses)\n",
    "        batch_sizes = torch.tensor(batch_sizes, device=valid_loss.device)\n",
    "        batch_size = torch.sum(batch_sizes)\n",
    "        valid_loss = torch.sum(valid_loss * batch_sizes) / batch_size\n",
    "\n",
    "        if torch.isnan(valid_loss):\n",
    "            raise Exception('Loss is NaN, training stopped.')\n",
    "\n",
    "        valid_loss_log = valid_loss.detach()\n",
    "        self.log(\n",
    "            'valid_loss',\n",
    "            valid_loss_log.item(),\n",
    "            batch_size=batch_size,\n",
    "            prog_bar=True,\n",
    "            on_epoch=True,\n",
    "        )\n",
    "        self.validation_step_outputs.append(valid_loss_log)\n",
    "        return valid_loss\n",
    "\n",
    "    def predict_step(self, batch, batch_idx):\n",
    "        if self.RECURRENT:\n",
    "            self.input_size = self.inference_input_size\n",
    "\n",
    "        temporal_cols = batch['temporal_cols']\n",
    "        windows_temporal, static, static_cols = self._create_windows(batch, step='predict')\n",
    "        n_windows = len(windows_temporal)\n",
    "        y_idx = batch['y_idx']\n",
    "\n",
    "        # Number of windows in batch\n",
    "        windows_batch_size = self.inference_windows_batch_size\n",
    "        if windows_batch_size < 0:\n",
    "            windows_batch_size = n_windows\n",
    "        n_batches = int(np.ceil(n_windows / windows_batch_size))\n",
    "        y_hats = []\n",
    "        for i in range(n_batches):\n",
    "            # Create and normalize windows [Ws, L+H, C]\n",
    "            w_idxs = np.arange(i*windows_batch_size, \n",
    "                    min((i+1)*windows_batch_size, n_windows))\n",
    "            windows = self._sample_windows(windows_temporal, static, static_cols, temporal_cols, step='predict', w_idxs=w_idxs)\n",
    "            windows = self._normalization(windows=windows, y_idx=y_idx)\n",
    "\n",
    "            # Parse windows\n",
    "            insample_y, insample_mask, _, _, \\\n",
    "                hist_exog, futr_exog, stat_exog = self._parse_windows(batch, windows)\n",
    "\n",
    "            if self.RECURRENT:                \n",
    "                y_hat = self._predict_step_recurrent_batch(insample_y=insample_y,\n",
    "                                                           insample_mask=insample_mask,\n",
    "                                                           futr_exog=futr_exog,\n",
    "                                                           hist_exog=hist_exog,\n",
    "                                                           stat_exog=stat_exog,\n",
    "                                                           y_idx=y_idx)\n",
    "            else:\n",
    "                y_hat = self._predict_step_direct_batch(insample_y=insample_y,\n",
    "                                                           insample_mask=insample_mask,\n",
    "                                                           futr_exog=futr_exog,\n",
    "                                                           hist_exog=hist_exog,\n",
    "                                                           stat_exog=stat_exog,\n",
    "                                                           y_idx=y_idx)                \n",
    "\n",
    "\n",
    "            y_hats.append(y_hat)\n",
    "        y_hat = torch.cat(y_hats, dim=0)\n",
    "        self.input_size = self.input_size_backup\n",
    "\n",
    "        return y_hat\n",
    "    \n",
    "    def fit(self, dataset, val_size=0, test_size=0, random_seed=None, distributed_config=None):\n",
    "        \"\"\" Fit.\n",
    "\n",
    "        The `fit` method, optimizes the neural network's weights using the\n",
    "        initialization parameters (`learning_rate`, `windows_batch_size`, ...)\n",
    "        and the `loss` function as defined during the initialization. \n",
    "        Within `fit` we use a PyTorch Lightning `Trainer` that\n",
    "        inherits the initialization's `self.trainer_kwargs`, to customize\n",
    "        its inputs, see [PL's trainer arguments](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.trainer.trainer.Trainer.html?highlight=trainer).\n",
    "\n",
    "        The method is designed to be compatible with SKLearn-like classes\n",
    "        and in particular to be compatible with the StatsForecast library.\n",
    "\n",
    "        By default the `model` is not saving training checkpoints to protect \n",
    "        disk memory, to get them change `enable_checkpointing=True` in `__init__`.\n",
    "\n",
    "        **Parameters:**<br>\n",
    "        `dataset`: NeuralForecast's `TimeSeriesDataset`, see [documentation](https://nixtla.github.io/neuralforecast/tsdataset.html).<br>\n",
    "        `val_size`: int, validation size for temporal cross-validation.<br>\n",
    "        `random_seed`: int=None, random_seed for pytorch initializer and numpy generators, overwrites model.__init__'s.<br>\n",
    "        `test_size`: int, test size for temporal cross-validation.<br>\n",
    "        \"\"\"\n",
    "        return self._fit(\n",
    "            dataset=dataset,\n",
    "            batch_size=self.batch_size,\n",
    "            valid_batch_size=self.valid_batch_size,\n",
    "            val_size=val_size,\n",
    "            test_size=test_size,\n",
    "            random_seed=random_seed,\n",
    "            distributed_config=distributed_config,\n",
    "        )\n",
    "\n",
    "    def predict(self, dataset, test_size=None, step_size=1,\n",
    "                random_seed=None, quantiles=None, **data_module_kwargs):\n",
    "        \"\"\" Predict.\n",
    "\n",
    "        Neural network prediction with PL's `Trainer` execution of `predict_step`.\n",
    "\n",
    "        **Parameters:**<br>\n",
    "        `dataset`: NeuralForecast's `TimeSeriesDataset`, see [documentation](https://nixtla.github.io/neuralforecast/tsdataset.html).<br>\n",
    "        `test_size`: int=None, test size for temporal cross-validation.<br>\n",
    "        `step_size`: int=1, Step size between each window.<br>\n",
    "        `random_seed`: int=None, random_seed for pytorch initializer and numpy generators, overwrites model.__init__'s.<br>\n",
    "        `quantiles`: list of floats, optional (default=None), target quantiles to predict. <br>\n",
    "        `**data_module_kwargs`: PL's TimeSeriesDataModule args, see [documentation](https://pytorch-lightning.readthedocs.io/en/1.6.1/extensions/datamodules.html#using-a-datamodule).\n",
    "        \"\"\"\n",
    "        self._check_exog(dataset)\n",
    "        self._restart_seed(random_seed)\n",
    "        if \"quantile\" in data_module_kwargs:\n",
    "            warnings.warn(\"The 'quantile' argument will be deprecated, use 'quantiles' instead.\")\n",
    "            if quantiles is not None:\n",
    "                raise ValueError(\"You can't specify quantile and quantiles.\")\n",
    "            quantiles = [data_module_kwargs.pop(\"quantile\")]\n",
    "        self._set_quantiles(quantiles)\n",
    "\n",
    "        self.predict_step_size = step_size\n",
    "        self.decompose_forecast = False\n",
    "        datamodule = TimeSeriesDataModule(dataset=dataset,\n",
    "                                          valid_batch_size=self.valid_batch_size,\n",
    "                                          **data_module_kwargs)\n",
    "\n",
    "        # Protect when case of multiple gpu. PL does not support return preds with multiple gpu.\n",
    "        pred_trainer_kwargs = self.trainer_kwargs.copy()\n",
    "        if (pred_trainer_kwargs.get('accelerator', None) == \"gpu\") and (torch.cuda.device_count() > 1):\n",
    "            pred_trainer_kwargs['devices'] = [0]\n",
    "\n",
    "        trainer = pl.Trainer(**pred_trainer_kwargs)\n",
    "        fcsts = trainer.predict(self, datamodule=datamodule)        \n",
    "        fcsts = torch.vstack(fcsts)\n",
    "\n",
    "        if self.MULTIVARIATE:\n",
    "            # [B, h, n_series (, Q)] -> [n_series, B, h (, Q)]\n",
    "            fcsts = fcsts.swapaxes(0, 2)\n",
    "            fcsts = fcsts.swapaxes(1, 2)\n",
    "\n",
    "        fcsts = tensor_to_numpy(fcsts).flatten()\n",
    "        fcsts = fcsts.reshape(-1, len(self.loss.output_names))\n",
    "        return fcsts\n",
    "\n",
    "    def decompose(self, dataset, step_size=1, random_seed=None, quantiles=None, **data_module_kwargs):\n",
    "        \"\"\" Decompose Predictions.\n",
    "\n",
    "        Decompose the predictions through the network's layers.\n",
    "        Available methods are `ESRNN`, `NHITS`, `NBEATS`, and `NBEATSx`.\n",
    "\n",
    "        **Parameters:**<br>\n",
    "        `dataset`: NeuralForecast's `TimeSeriesDataset`, see [documentation here](https://nixtla.github.io/neuralforecast/tsdataset.html).<br>\n",
    "        `step_size`: int=1, step size between each window of temporal data.<br>\n",
    "        `quantiles`: list of floats, optional (default=None), target quantiles to predict. <br>\n",
    "        `**data_module_kwargs`: PL's TimeSeriesDataModule args, see [documentation](https://pytorch-lightning.readthedocs.io/en/1.6.1/extensions/datamodules.html#using-a-datamodule).\n",
    "        \"\"\"\n",
    "        # Restart random seed\n",
    "        if random_seed is None:\n",
    "            random_seed = self.random_seed\n",
    "        torch.manual_seed(random_seed)\n",
    "        self._set_quantiles(quantiles)\n",
    "\n",
    "        self.predict_step_size = step_size\n",
    "        self.decompose_forecast = True\n",
    "        datamodule = TimeSeriesDataModule(dataset=dataset,\n",
    "                                          valid_batch_size=self.valid_batch_size,\n",
    "                                          **data_module_kwargs)\n",
    "        trainer = pl.Trainer(**self.trainer_kwargs)\n",
    "        fcsts = trainer.predict(self, datamodule=datamodule)\n",
    "        self.decompose_forecast = False # Default decomposition back to false\n",
    "        fcsts = torch.vstack(fcsts)\n",
    "        return tensor_to_numpy(fcsts)        "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
